"""
LLM (大语言模型) - 核心控制器
这是整个AI系统的大脑，负责协调各个组件工作
"""

import numpy as np
from typing import List, Protocol, runtime_checkable
from vocab import Vocab

# 常量定义 - 就像设定游戏规则
MAX_SEQ_LEN = 40      # 最大句子长度（一次最多处理40个词）
EMBEDDING_DIM = 32    # 词向量维度（每个词用32个数字表示）
HIDDEN_DIM = 32       # 隐藏层维度（内部处理时的数据宽度）

@runtime_checkable
class Layer(Protocol):
    """
    层接口协议 - 就像定义所有AI组件必须遵守的规则
    每个AI组件（注意力层、前馈层等）都必须实现这些方法
    """
    def layer_type(self) -> str:
        """返回层的类型名称，比如'注意力层'、'前馈层'等"""
        ...
    
    def forward(self, input_data: np.ndarray) -> np.ndarray:
        """前向传播：处理输入数据，返回处理结果"""
        ...
    
    def backward(self, grads: np.ndarray, lr: float) -> np.ndarray:
        """反向传播：根据错误信号调整参数，学习改进"""
        ...

class LLM:
    """
    大语言模型主类 - AI的大脑总控制器
    就像一个智能助手的大脑，能理解问题并生成回答
    """
    
    def __init__(self, vocab: Vocab, network: List[Layer]):
        """
        初始化AI大脑
        vocab: 词汇表（AI认识的所有词汇）
        network: 神经网络层列表（AI的思考组件）
        """
        self.vocab = vocab          # AI的词汇表，就像字典
        self.network = network      # AI的神经网络，就像大脑的各个区域
    
    def network_description(self) -> str:
        """
        描述AI的网络结构
        返回：网络各层的名称，比如"词嵌入层, 注意力层, 输出层"
        """
        return ", ".join([layer.layer_type() for layer in self.network])
    
    def predict(self, text: str) -> str:
        """
        AI预测/生成回答的主入口
        输入：用户的问题文本
        输出：AI生成的回答文本
        
        就像问AI"天空是什么颜色？"，AI回答"蓝色"
        """
        # 第1步：让AI的大脑处理问题，得到词汇ID列表
        output_tokens = self.forward(text)
        
        # 第2步：检查是否有输出
        if not output_tokens:
            return ""  # 如果AI没想出什么，就返回空字符串
        
        # 第3步：把词汇ID转换回文字
        token_strs = []
        for token_id in output_tokens:
            if token_id in self.vocab.decode:  # 确保这个ID在词汇表中
                token_strs.append(self.vocab.decode[token_id])
        
        # 第4步：把词汇组合成句子返回
        return " ".join(token_strs)
    
    def forward(self, text: str) -> List[int]:
        """
        前向传播 - AI思考过程的核心
        输入：文本字符串
        输出：生成的词汇ID列表
        
        就像AI读懂问题后，一个词一个词地组织回答
        """
        # 第1步：把文本切分成词汇ID
        tokenized = self.tokenize(text)
        output_tokens: List[int] = []  # 存储AI生成的词汇
        
        # 安全检查：确保至少有一个词
        if not tokenized:
            return output_tokens
        
        input_len = len(tokenized)
        
        # 防止输入太长超出处理能力
        if input_len >= MAX_SEQ_LEN:
            return output_tokens
        
        # 开始生成回答，一次生成一个词
        for _ in range(MAX_SEQ_LEN - input_len):
            # 检查是否已经生成足够多的词
            if len(output_tokens) >= MAX_SEQ_LEN - 1:
                break
            
            # 第2步：把当前的词汇序列转换成数字矩阵
            token_input = np.array(tokenized, dtype=np.float32).reshape(1, -1)
            input_data = token_input
            
            # 第3步：让数据通过AI的各个思考层
            for layer in self.network:
                input_data = layer.forward(input_data)  # 每一层都处理一下数据
            
            logits = input_data  # 最终得到的是每个词的"得分"
            
            # 安全检查
            if logits.shape[0] == 0:
                break
            
            # 第4步：取最后一个位置的得分（这是对下一个词的预测）
            last_logit = logits[-1:, :]
            
            # 第5步：把得分转换成概率分布（softmax）
            # 就像把考试分数转换成百分比
            probs = self.softmax(last_logit)  # 形状: 1 x 词汇表大小
            
            # 第6步：选择概率最高的词（贪心解码）
            # 就像选择最可能的答案
            tokens = self.greedy_decode(probs)
            next_token = tokens[-1]  # 获取选中的词
            
            # 第7步：把新词加入到输出和输入中
            output_tokens.append(next_token)
            tokenized.append(next_token)
            
            # 第8步：如果遇到结束符，停止生成
            if next_token == self.vocab.encode_word("</s>"):
                break
        
        return output_tokens
    
    def train(self, data: List[str], epochs: int, lr: float):
        """
        训练AI - 让AI学习如何更好地回答问题
        data: 训练数据（问题和答案的例子）
        epochs: 训练轮数（学习多少遍）
        lr: 学习率（学习的步长大小）
        
        就像给学生做练习题，做得越多越熟练
        """
        print(f"开始训练AI，共{len(data)}条数据，训练{epochs}轮")
        
        # 第1步：把所有训练文本转换成词汇ID
        tokenized_data = []
        for input_text in data:
            tokens = self.tokenize(input_text)
            tokenized_data.append(tokens)
        
        # 第2步：开始训练循环
        for epoch in range(epochs):
            total_loss = 0.0  # 记录这一轮的总错误
            
            # 第3步：用每条训练数据训练一次
            for training_row in tokenized_data:
                if len(training_row) < 2:  # 太短的句子跳过
                    continue
                
                # 第4步：准备输入和目标
                # 输入：除了最后一个词的所有词
                # 目标：除了第一个词的所有词
                # 这样AI学习"看到前面的词，预测下一个词"
                input_ids = training_row[:-1]   # [我, 爱, 北] 
                target_ids = training_row[1:]   # [爱, 北, 京]
                
                # 第5步：前向传播 - AI尝试预测
                input_data = np.zeros((1, len(input_ids)), dtype=np.float32)
                input_data[0, :] = input_ids
                
                # 让数据通过AI的各层
                for layer in self.network:
                    input_data = layer.forward(input_data)
                
                logits = input_data  # AI的预测得分
                probs = self.softmax(logits)  # 转换成概率
                
                # 第6步：计算错误（损失）
                # 比较AI的预测和正确答案，看差距有多大
                loss = self.cross_entropy_loss_step(probs, target_ids)
                total_loss += loss
                
                # 第7步：反向传播 - 根据错误调整AI参数
                # 计算梯度（错误的方向和大小）
                grads_output = self.compute_gradients_step(probs, target_ids)
                
                # 梯度裁剪：防止学习步长太大
                self.clip_gradients(grads_output, 5.0)
                
                # 从后往前调整每一层的参数
                for layer in reversed(self.network):
                    grads_output = layer.backward(grads_output, lr)
                
                # 检查预测结果
                tokens = self.greedy_decode(probs)
                next_token = tokens[-1]
                
                # 如果预测到结束符，跳过这个样本
                if next_token == self.vocab.encode_word("</s>"):
                    continue
            
            # 第8步：打印这一轮的平均错误
            avg_loss = total_loss / len(tokenized_data)
            print(f"第{epoch}轮训练: 平均损失 = {avg_loss:.4f}")
    
    def tokenize(self, text: str) -> List[int]:
        """
        分词器 - 把文本切分成词汇，并转换成ID
        输入："你好吗？"
        输出：[1, 2, 3] （对应词汇表中的ID）
        
        就像把句子切分成一个个单词
        """
        tokens = []  # 存储词汇ID
        
        # 按空格分割文本
        for word in text.split():
            # 特殊处理结束符
            if word == "</s>":
                token_id = self.vocab.encode_word(word)
                if token_id is not None:
                    tokens.append(token_id)
                continue
            
            current_word = ""  # 当前正在处理的词
            
            # 逐字符处理，分离标点符号
            for c in word:
                if c in ".,!?;:":  # 如果是标点符号
                    # 先保存之前的词
                    if current_word:
                        token_id = self.vocab.encode_word(current_word)
                        if token_id is not None:
                            tokens.append(token_id)
                        current_word = ""
                    
                    # 标点符号单独作为一个词
                    token_id = self.vocab.encode_word(c)
                    if token_id is not None:
                        tokens.append(token_id)
                else:
                    current_word += c  # 累积字符形成词
            
            # 保存最后的词
            if current_word:
                token_id = self.vocab.encode_word(current_word)
                if token_id is not None:
                    tokens.append(token_id)
        
        return tokens
    
    @staticmethod
    def softmax(logits: np.ndarray) -> np.ndarray:
        """
        Softmax函数 - 把得分转换成概率分布
        输入：[3.2, 1.8, 0.5] （各个词的得分）
        输出：[0.7, 0.2, 0.1] （各个词的概率，总和为1）
        
        就像把考试分数转换成百分比排名
        """
        result = logits.copy()
        
        # 对每一行应用softmax
        for i in range(result.shape[0]):
            row = result[i, :]
            
            # 数值稳定性技巧：减去最大值防止溢出
            max_val = np.max(row)
            exp_values = np.exp(row - max_val)  # 计算指数
            sum_exp = np.sum(exp_values)        # 求和
            
            # 归一化得到概率
            result[i, :] = exp_values / sum_exp
        
        return result
    
    @staticmethod
    def greedy_decode(probs: np.ndarray) -> List[int]:
        """
        贪心解码 - 选择概率最高的词
        输入：[[0.1, 0.7, 0.2]] （每个词的概率）
        输出：[1] （概率最高的词的索引）
        
        就像选择题中选择最可能正确的答案
        """
        return [np.argmax(row) for row in probs]
    
    @staticmethod
    def cross_entropy_loss_step(probs: np.ndarray, target: List[int]) -> float:
        """
        交叉熵损失 - 计算AI预测和正确答案之间的差距
        probs: AI预测的概率分布
        target: 正确答案的词汇ID
        
        返回值越小说明预测越准确，越大说明差距越大
        就像计算考试错误的严重程度
        """
        loss = 0.0
        
        for row_idx in range(probs.shape[0]):
            # 获取正确答案对应的预测概率
            prob_target = probs[row_idx, target[row_idx]]
            
            # 计算负对数似然（数值稳定版本）
            loss -= np.log(max(prob_target, 1e-15))  # 防止log(0)
        
        return loss / len(target)  # 返回平均损失
    
    @staticmethod
    def compute_gradients_step(probs: np.ndarray, target: List[int]) -> np.ndarray:
        """
        计算梯度 - 确定参数应该如何调整
        probs: AI的预测概率
        target: 正确答案
        
        返回梯度矩阵，指示每个参数应该增大还是减小
        就像告诉学生哪些地方答错了，应该怎么改正
        """
        grads = probs.copy()  # 从预测概率开始
        
        # 检查维度匹配
        if probs.shape[0] != len(target):
            raise ValueError("预测结果和目标答案的数量不匹配")
        
        batch_size = len(target)
        
        # 计算softmax + 交叉熵的梯度：预测概率 - 真实标签
        for row_idx in range(grads.shape[0]):
            grads[row_idx, target[row_idx]] -= 1.0  # 正确答案位置减1
        
        # 按批次大小归一化，保持训练稳定
        grads = grads / batch_size
        
        return grads
    
    @staticmethod
    def clip_gradients(grads: np.ndarray, max_norm: float):
        """
        梯度裁剪 - 防止学习步长过大
        grads: 梯度矩阵
        max_norm: 允许的最大梯度范数
        
        就像限制学生改正错误时不要用力过猛
        """
        # 计算梯度的L2范数（总体大小）
        norm = np.sqrt(np.sum(grads ** 2))
        
        # 如果梯度太大，按比例缩小
        if norm > max_norm:
            scale = max_norm / norm
            grads *= scale  # 直接修改原数组